package com.xfsg.core.kafka;

import com.yomahub.tlog.constant.TLogConstants;
import com.yomahub.tlog.core.rpc.TLogLabelBean;
import com.yomahub.tlog.core.rpc.TLogRPCHandler;
import lombok.extern.slf4j.Slf4j;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.common.header.Headers;
import org.springframework.kafka.listener.MessageListener;

/**
 * @author ZhangYi
 * @version 1.0.0
 * @Description:
 * @date 2022/3/28
 */
@Slf4j
public class TracingMessageListenerMethodInterceptor<T extends MessageListener> implements MethodInterceptor {

    private static TLogRPCHandler tLogRPCHandler = new TLogRPCHandler();

    @Override
    public Object invoke(MethodInvocation invocation) throws Throwable {
        if (!"onMessage".equals(invocation.getMethod().getName())) {
            return invocation.proceed();
        }
        Object[] arguments = invocation.getArguments();
        Object record = record(arguments);
        if (record == null) {
            return invocation.proceed();
        }
        Headers headers = ((ConsumerRecord<?, ?>) record).headers();
        if (headers.toArray().length > 0){
            String traceId = KafkaHeaderUtil.lastStringHeader(headers,TLogConstants.TLOG_TRACE_KEY);
            String preIvkApp = KafkaHeaderUtil.lastStringHeader(headers,TLogConstants.PRE_IVK_APP_KEY);
            String preIvkHost = KafkaHeaderUtil.lastStringHeader(headers,TLogConstants.PRE_IVK_APP_HOST);
            String preIp = KafkaHeaderUtil.lastStringHeader(headers,TLogConstants.PRE_IP_KEY);
            String spanId = KafkaHeaderUtil.lastStringHeader(headers,TLogConstants.TLOG_SPANID_KEY);

            TLogLabelBean labelBean = new TLogLabelBean(preIvkApp, preIvkHost, preIp, traceId, spanId);
            tLogRPCHandler.processProviderSide(labelBean);
        }

        if (log.isDebugEnabled()) {
            log.debug("Wrapping onMessage call");
        }
        try  {
            return invocation.proceed();
        }
        catch (RuntimeException | Error e) {
            String message = e.getMessage();
            if (message == null) {
                message = e.getClass().getSimpleName();
            }
            throw e;
        }
        finally {
            tLogRPCHandler.cleanThreadLocal();
        }
    }

    private Object record(Object[] arguments) {
        for (Object object : arguments) {
            if (object instanceof ConsumerRecord) {
                return object;
            }
        }
        return null;
    }
}
